import pygame
import socket
import json
import threading
import time

# 初始化 PyGame
pygame.init()

# 游戏窗口设置
WINDOW_WIDTH = 800
WINDOW_HEIGHT = 600
window = pygame.display.set_mode((WINDOW_WIDTH, WINDOW_HEIGHT))
pygame.display.set_caption("UDP 游戏客户端")

# 颜色定义
WHITE = (255, 255, 255)
BLACK = (0, 0, 0)
RED = (255, 0, 0)
GREEN = (0, 255, 0)
BLUE = (0, 0, 255)

# 服务器设置
SERVER_ADDRESS = ('qsont.xyz', 8002)

# 玩家设置
player_id = None
players = {}  # 存储其他玩家信息 {id: {'x': x, 'y': y}}

# 创建 UDP 套接字
sock = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
sock.settimeout(1.0)  # 设置超时时间

# 发送心跳的函数
def send_heartbeat():
    while True:
        try:
            if player_id:  # 确保我们已经有了玩家ID
                heartbeat_msg = {
                    'type': 'heartbeat',
                    'id': player_id
                }
                sock.sendto(json.dumps(heartbeat_msg).encode(), SERVER_ADDRESS)
        except Exception as e:
            print(f"发送心跳失败: {e}")
        time.sleep(1)  # 每秒发送一次心跳

# 接收服务器消息的函数
def receive_messages():
    global player_id
    while True:
        try:
            data, addr = sock.recvfrom(1024)
            message = json.loads(data.decode())
            
            if message['type'] == 'welcome':
                # 服务器欢迎消息，包含我们的玩家ID
                player_id = message['id']
                print(f"连接到服务器，玩家ID: {player_id}")
                
            elif message['type'] == 'update':
                # 更新其他玩家位置
                for pid, pos in message['players'].items():
                    players[pid] = pos
                    
        except socket.timeout:
            continue
        except Exception as e:
            print(f"接收消息错误: {e}")

# 启动心跳线程
heartbeat_thread = threading.Thread(target=send_heartbeat)
heartbeat_thread.daemon = True
heartbeat_thread.start()

# 启动接收消息线程
receive_thread = threading.Thread(target=receive_messages)
receive_thread.daemon = True
receive_thread.start()

# 发送初始连接消息
try:
    connect_msg = {'type': 'connect'}
    sock.sendto(json.dumps(connect_msg).encode(), SERVER_ADDRESS)
except Exception as e:
    print(f"初始连接失败: {e}")

# 游戏主循环
clock = pygame.time.Clock()
running = True

while running:
    # 处理事件
    for event in pygame.event.get():
        if event.type == pygame.QUIT:
            running = False
    
    # 处理键盘输入
    keys = pygame.key.get_pressed()
    move_msg = {'type': 'move', 'id': player_id, 'dx': 0, 'dy': 0}
    
    if keys[pygame.K_LEFT]:
        move_msg['dx'] = -10
    if keys[pygame.K_RIGHT]:
        move_msg['dx'] = 10
    if keys[pygame.K_UP]:
        move_msg['dy'] = -10
    if keys[pygame.K_DOWN]:
        move_msg['dy'] = 10
    
    # 如果有移动，发送移动消息
    if move_msg['dx'] != 0 or move_msg['dy'] != 0:
        try:
            sock.sendto(json.dumps(move_msg).encode(), SERVER_ADDRESS)
        except Exception as e:
            print(f"发送移动消息失败: {e}")
    
    # 绘制游戏
    window.fill(BLACK)
    
    # 绘制所有玩家
    for pid, pos in players.items():
        color = GREEN if pid == player_id else RED  # 自己显示为绿色，其他玩家显示为红色
        pygame.draw.circle(window, color, (pos['x'], pos['y']), 20)
        # 显示玩家ID
        font = pygame.font.SysFont(None, 24)
        text = font.render(str(pid), True, WHITE)
        window.blit(text, (pos['x'] - 10, pos['y'] - 30))
    
    pygame.display.flip()
    clock.tick(60)  # 60 FPS

# 退出游戏
pygame.quit()
sock.close()
